'''
pyrwi.daq - Helper functions for interacting with a DAQ.

Currently interfaces with the nidaqmx module
'''
import scipy as sp
import nidaqmx as DAQmx
from collections import deque

class Enum :
    '''
    Generate Enumerated types.
    Basically assign unique integers to a set of names
    '''
    def __init__(self) :
        self.nxt_int = 0 #Next item in enum will get this number
        self.used_int = {} #Store a list of integers already assigned
        #TODO: Maybe a sorted list would be better?
    
    def new(self, int=None, name="none") :
        if int is None :
            rtn = self.nxt_int
            self.nxt_int += 1
        else :
            #Check self.used_int for conflict
            if self.used_int.has_key(int) :
                raise SyntaxError("$d is already used in enum" % int)
            rtn = int
        
        self.used_int[rtn] = name #Store integer as used
        #Make sure I am ready for the next call to self.new()
        while self.used_int.has_key(self.nxt_int) :
            self.nxt_int += 1
        
        return rtn
        
DataTypes = Enum()
float64 = DataTypes.new("float64")
int16 = DataTypes.new("int16")
int32 = DataTypes.new("int32")
uInt8 = DataTypes.new("uInt8")
uInt16 = DataTypes.new("uInt16")
uInt32 = DataTypes.new("uInt32")

formatStrings = { float64: "%g", int16: "%d", int32: "%d",
                    uInt8: "%d", uInt16: "%d", uInt32: "%d", }

def ChannelDataType(taskHandle) :
    '''
    Pick a data type to use given the current task configuration.
    Not the following assumptions
    Analog is always float64
    Digital is always uInt32
    Counters are always uInt32
    '''
    type = DAQmx.GetChanType(taskHandle, "")
    if type == DAQmx.Val_AI or type == DAQmx.Val_AO : return float64
    if type == DAQmx.Val_DI or type == DAQmx.Val_DO : return uInt32
    if type == DAQmx.Val_CI or type == DAQmx.Val_CO : return uInt32
    
    raise TypeError("Uknown channel type: %d" % type)

class ReadChannelHelper :
    """
    Help handle the common tasks of reading, logging, and plotting
    a DAQ read task.
    
    autoWrite - Default: False
    """
    def __init__(self, taskHandle, plot=None, logFile=None, buf_time=0) :
        '''
        Create a ReadChannelHelper
        NOTE: taskHandle must already be configured, because a great deal
                of configuration is loaded from the task at object creation
        
        @param taskHandle: DAQmx task handle
        @param plot: FL_Plot widget
        @param logFile: File that data will be written to. File handle or path
        @buf_time: How many seconds to store in buffer, 0 for unlimited
        '''
        #--- Widgets and handles
        self.taskH = taskHandle
        self.plot = plot
        if logFile is None :
            self.file = None
        elif isinstance(logFile, file) : 
            #Already a file handle
            self.file = logFile
        else :
            #Open a file, default to overwriting any existing file
            self.file = open(logFile, 'w')
        
        #--- Configurations
        self.chan_cnt = DAQmx.GetReadNumChans(self.taskH)
        self.Fs = DAQmx.GetSampClkRate(self.taskH) # Samples / sec / channel
        self.dataType = ChannelDataType(self.taskH)
        self.fmt_str = formatStrings[self.dataType] # Number format for % operator

        self.bufferMaxLen = int(buf_time * self.Fs) #Max length of buffer, 0 for unlimited

        #--- Initialize State
        self.clear()
                    
    def clear(self) :
        '''
        Clear the buffer and the plot.
        TODO: What to do with the file?
        '''
        self.x = None
        #self.buffer = [[]]*self.chan_cnt #Data from DAQ one list for ea channel
        self.buffer = []
        for i in xrange(self.chan_cnt) :
            if self.bufferMaxLen == 0 :
                self.buffer.append( deque([]) )
            else :
                self.buffer.append( deque([], self.bufferMaxLen) )
            
        self.buffer_len = 0 #Actual length of buffer
        self.t0 = 0 #Time of the self.data[0] point
        self.iSaved = 0 #All points prior to this time have been saved
        
        #Build an X axis
        self.x = sp.arange(self.t0,stop=self.t0+float(self.bufferMaxLen)/self.Fs, step=1.0/self.Fs)
        
        if self.plot : self.plot.clear()

    def save(self, elemShifted=None) :
        """
        Write the buffer contents to a file.
        """
        if self.file is None : return #No file to save to
        if elemShifted is not None and self.iSaved > elemShifted : 
            #Data being shifted out has already been saved
            return

        
        if self.bufferMaxLen == 0 and len(self.buffer)>0:
            #Unlimited buffers have no x data, so we must create it.
            buf_len = len(self.buffer[0])
            #Build an X axis
            self.x = sp.arange(self.t0,stop=self.t0+float(buf_len)/self.Fs, step=1.0/self.Fs)

        pat = "%f, " + ",".join([self.fmt_str]*len(self.buffer)) + "\n"
        
        for i,tpl in enumerate(zip(self.x[0:self.buffer_len], *self.buffer)) :
            if i > self.iSaved :
                self.file.write( pat % tpl )
                        
        #Store the last time saved to file
        self.iSaved = len(self.x)-1
            
    def flush(self) :
        '''
        Make sure entire buffer is written to file
        '''
        self.save()
        if self.file is not None: self.file.flush()

    def ReadBuffer(self, nSamples) :
        read = DAQmx.ReadAnalogF64(self.taskH, nSamples, 10.0, \
            DAQmx.Val_GroupByChannel, nSamples*self.chan_cnt)
        self.cb_AnalogRead(nSamples, read)

    def cb_AnalogRead(self, sampsPerChanRead, dataRead) :
        """
        Call back to be supplied to DAQmx.RegisterEveryNSamplesEvent when
        DAQmx.ReadAnalogF64_cb is used
        """
        self.addToBuffer(sampsPerChanRead, dataRead)
        self.update()

    def cb_DigitalRead(self, sampsPerChanRead, dataRead) :
        """
        Call back to be supplied to DAQmx.RegisterEveryNSamplesEvent when
        DAQmx.ReadDigitalU32_cb is used
        """
        #TODO
        self.addToBuffer(sampsPerChanRead, dataRead)
        self.update()

    def addToBuffer(self, sampsPerChan, data) :
        '''
        Add new data to the buffer.
        @param sampsPerChan: Number of samples read for each channel
        @param data: Array returned by DAQ Read, always in 
                        DAQmx.Val_GroupByChannel format
        '''
        if self.bufferMaxLen != 0 :            
            #How many samples do I need to shift out?
            shiftData = sampsPerChan + self.buffer_len - self.bufferMaxLen
            if shiftData < 0 : shiftData = 0 #No negative shifts
            
            if shiftData > 0 :
                #Update starting time
                self.t0 += float(shiftData) / self.Fs
                self.save(shiftData) #Store data before I erase it
                
                #Shift time in x to match buffer
                self.x += self.t0 - self.x[0]
                self.iSaved -= shiftData
        
        for i in xrange(self.chan_cnt) :
            #If MaxLen is set, the deque will remove enough samples so 
            #these can be added
            self.buffer[i].extend( data[i*sampsPerChan:(i+1)*sampsPerChan] )
        
        self.buffer_len =  len(self.buffer[0])
                    
    def update(self) :
        '''
        Update all the attached elements (plot, maybe file)
        '''
        if self.plot is not None :
            if len(self.plot.lines) == 0 :
                #Plot widget is empty, so plot data for the first time
                for d in self.buffer :
                    self.plot.plot(self.x[:self.buffer_len],d)
            else :
                self.plot.update(self.x[:self.buffer_len],*self.buffer)

        #TODO: File
    def close(self) :
        '''
        Clean up
        '''
        self.save()
        if self.file is not None: self.file.close()
        pass

if __name__ == "__main__":
    import time
    nSamples = 10
    
    task = DAQmx.CreateTask("ContAcqTest")             #Create a new task
    DAQmx.CreateAIVoltageChan(task,"Dev1/ai0:1","",\
            DAQmx.Val_Cfg_Default,-10.0,10.0,\
            DAQmx.Val_Volts,None)                       #Configure to read a voltage

    DAQmx.CfgSampClkTiming(task,"",5*nSamples,\
            DAQmx.Val_Rising,DAQmx.Val_ContSamps,\
            nSamples) #Configure a continuous clock
            
    readHelper = ReadChannelHelper(task, None, logFile='test.txt', buf_time=1.0)
    
    DAQmx.RegisterEveryNSamplesEvent(task, DAQmx.Val_Acquired_Into_Buffer,\
            nSamples,0,DAQmx.ReadAnalogF64_cb, readHelper.cb_AnalogRead)
    
    if True :
        import timeit
        data = sp.empty(2*nSamples)
        data[:nSamples] = sp.rand(nSamples)*0.01 + 0.64
        data[nSamples:] = sp.rand(nSamples)*0.01 + 0.7
        print len(data)
        
        def addData() :
            readHelper.clear()
            for i in xrange(50) :
                readHelper.addToBuffer(nSamples, data)
                readHelper.update()
            return
            
        print timeit.timeit('addData()', "from __main__ import addData", number=50)
        #addData()
        readHelper.close()
        
    else :    
        DAQmx.StartTask(task)
        time.sleep(2)

    DAQmx.StopTask(task)
    DAQmx.ClearTask(task)

